# secure_config.py (配置安全模块)
import os
from cryptography.fernet import Fernet
import yaml
from yaml import SafeLoader


class SecureConfigLoader:
    _fernet_key = os.getenv("CONFIG_KEY")
    if not _fernet_key:
        raise ValueError("请设置环境变量 CONFIG_KEY")
    _fernet = Fernet(_fernet_key.encode())

    @classmethod
    def _decrypt(cls, ciphertext):
        """解密方法"""
        try:
            return cls._fernet.decrypt(ciphertext.encode()).decode()
        except Exception as e:
            raise ValueError(f"解密失败: {e}")

    @classmethod
    def load(cls, file_path):
        # 注册 !encrypt 标签处理器
        def encrypt_constructor(loader, node):
            value = loader.construct_scalar(node)
            return cls._decrypt_values(value)  # 调用解密方法

        # 添加标签处理器到 SafeLoader
        SafeLoader.add_constructor("!encrypt", encrypt_constructor)

        # 修复点：添加 encoding='utf-8'
        with open(file_path, "r", encoding="utf-8") as f:
            # 应该建议用户使用yaml.SafeLoader而不是FullLoader，以增强安全性，同时注册自定义标签处理器，确保只有允许的标签被处理，避免潜在的安全风险。
            config = yaml.load(f, Loader=yaml.SafeLoader)
        return cls._decrypt_values(config)

    @classmethod
    def _decrypt_values(cls, config):
        if isinstance(config, dict):
            return {k: cls._decrypt_values(v) for k, v in config.items()}
        elif isinstance(config, list):
            return [cls._decrypt_values(v) for v in config]
        elif isinstance(config, str) and config.startswith("!encrypt "):
            ciphertext = config.split(" ", 1)[1].encode()
            return cls._fernet.decrypt(ciphertext).decode()
        return config